package heartbeat;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.security.NoSuchAlgorithmException;
import java.util.LinkedList;
import java.util.Queue;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;

public class SslTcpConnection extends TcpConnection 
{
	private SSLContext mSslContext;
	private SSLEngine mEngine;
	private ByteBuffer mInNetData;
	private ByteBuffer mOutNetData;
	private ByteBuffer mInPeerData;
	private boolean mNeedToRead = false;
	private int mNetBufSize;
	private int mAppBufSize;
	protected Queue<Packet> mIncomingPackets;
	
	public SslTcpConnection(String address, int port) throws IOException
	{
		super(address, port);
		mConnected = false;
		
		try 
		{
			mSslContext = SSLContext.getDefault();
		} 
		catch (NoSuchAlgorithmException e) 
		{
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		mEngine = mSslContext.createSSLEngine();
		mEngine.setUseClientMode(true);
		mEngine.setEnabledCipherSuites(new String[]{"SSL_RSA_WITH_RC4_128_SHA"});
		mEngine.setEnabledProtocols(new String[]{"TLSv1"});
	    mNetBufSize = mEngine.getSession().getPacketBufferSize();
	    mInNetData = ByteBuffer.allocateDirect(mNetBufSize);
	    mOutNetData = ByteBuffer.allocateDirect(mNetBufSize);
	    mAppBufSize = mEngine.getSession().getApplicationBufferSize();
	    mInPeerData = ByteBuffer.allocateDirect(mAppBufSize);
	    mIncomingPackets = new LinkedList<Packet>();
		doHandshake();
	}
	
	protected void doHandshake() throws IOException
	{
		// Create byte buffers to use for holding application data
	    int appBufferSize = mEngine.getSession().getApplicationBufferSize();
	    ByteBuffer myAppData = ByteBuffer.allocate(appBufferSize);
	    ByteBuffer peerAppData = ByteBuffer.allocate(appBufferSize);

	    // Begin handshake
	    mEngine.beginHandshake();
	    SSLEngineResult.HandshakeStatus hs = mEngine.getHandshakeStatus();

	    // Process handshaking message
	    while (hs != SSLEngineResult.HandshakeStatus.FINISHED &&
	        hs != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) {

	        switch (hs) {

	        case NEED_UNWRAP:
	            // Receive handshaking data from peer
	        	int bytes = mSocketChannel.read(mInNetData);
	        	if(bytes < 0)
	            {
	            	System.out.println("Closed channel error");
	                // Handle closed channel
	            }
	        	
	        	System.out.println(bytes);

	            // Process incoming handshaking data
	            mInNetData.flip();
	            SSLEngineResult res = mEngine.unwrap(mInNetData, peerAppData);
	            System.out.println("Unwrapped");
	            System.out.println(res.getHandshakeStatus());
	            mInNetData.compact();
	            hs = res.getHandshakeStatus();

	            // Check status
	            switch (res.getStatus()) 
	            {
	            	case OK:
	                // Handle OK status
	            		break;
				case BUFFER_OVERFLOW:
					break;
				case BUFFER_UNDERFLOW:
					break;
				case CLOSED:
					break;
				default:
					break;

	            // Handle other status: BUFFER_UNDERFLOW, BUFFER_OVERFLOW, CLOSED
	            }
	            break;

	        case NEED_WRAP :
	            // Empty the local network packet buffer.
	            mOutNetData.clear();

	            // Generate handshaking data
	            res = mEngine.wrap(myAppData, mOutNetData);
	            hs = res.getHandshakeStatus();

	            // Check status
	            switch (res.getStatus()) {
	            case OK :
	                mOutNetData.flip();

	                // Send the handshaking data to peer
	                while (mOutNetData.hasRemaining()) 
	                {
	                    if (mSocketChannel.write(mOutNetData) < 0) 
	                    {
	                       throw new IOException("Bad write");
	                    }
	                }
	                break;
				case BUFFER_OVERFLOW:
					break;
				case BUFFER_UNDERFLOW:
					break;
				case CLOSED:
					break;
				default:
					break;
	            }
	            break;

	        case NEED_TASK :
	        	Runnable r = null;
	            while ( (r = mEngine.getDelegatedTask()) != null) {
	                r.run();
	            }
	            hs = mEngine.getHandshakeStatus();
	            break;

	        case FINISHED :
	        	mConnected = true;
	        // Handle other status:  // FINISHED or NOT_HANDSHAKING
			case NOT_HANDSHAKING:
				break;
			default:
				break;
	        }
	    }
	    
	    if(hs == SSLEngineResult.HandshakeStatus.FINISHED)
	    {
	    	mInNetData.clear();
	    	mNeedToRead = true;
	    	mConnected = true;
	    }
	}
	
	public int byteBufferCopy(ByteBuffer src, ByteBuffer dst)
	{
		src.flip();
		int maxTransfer = Math.min(src.remaining(), dst.remaining());
		ByteBuffer bufTmp = mInPeerData.duplicate();
		bufTmp.limit(bufTmp.position() + maxTransfer);
		dst.put(bufTmp);
		src.position(src.position() + maxTransfer);
		src.compact();
		return maxTransfer;
	}
	
	public int readInternal(ByteBuffer dst) throws IOException
	{
		int bytesRequested = dst.remaining();
		int bytesAvailable = mInPeerData.position();
		if(bytesAvailable >= bytesRequested)
		{
			int transferred = byteBufferCopy(mInPeerData, dst);
			return transferred;
		}
	
		if(!mSocketChannel.isConnected())
		{
			System.out.println("!mSocketChannel.isConnected");
			return 0;
		}

		if(!mConnected)
		{
			System.out.println("!mSocketChannel.isConnected");
			return 0;
		}

		if(mNeedToRead)
		{
			
			mNeedToRead = false;
			mSocketChannel.read(mInNetData);
			mInNetData.flip();
		}

		SSLEngineResult res = mEngine.unwrap(mInNetData, mInPeerData);
		if(res.getHandshakeStatus() != SSLEngineResult.HandshakeStatus.FINISHED && res.getHandshakeStatus() != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING)
		{
			System.out.println("HANDSHAKE STATUS BAD");
			return -1;
		}

		if(res.getStatus() == SSLEngineResult.Status.BUFFER_UNDERFLOW)
		{
			if(mEngine.getSession().getApplicationBufferSize() > mInPeerData.capacity())
				System.out.println("under1");
			
			mInNetData.compact();
			mNeedToRead = true;
			return 0;
		}
			
		return byteBufferCopy(mInPeerData, dst);
	}
	
	protected void writeInternal(ByteBuffer buf) throws IOException
	{
		ByteBuffer outBuf = ByteBuffer.allocate(mNetBufSize);
		outBuf.clear();
		// Generate SSL/TLS encoded data (handshake or application data)
		while(buf.hasRemaining())
		{
			SSLEngineResult res = mEngine.wrap(buf, outBuf);

			// Process status of call
			if (res.getStatus() == SSLEngineResult.Status.OK) 
			{
				outBuf.flip();
				// Send SSL/TLS encoded data to peer
				while(outBuf.hasRemaining()) 
				{
					int num = mSocketChannel.write(outBuf);
					if (num < 0) 
						throw new IOException("Bad write");
				}
				
				outBuf.clear();
			}
		}
	}
}
